{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc95e213",
   "metadata": {},
   "source": [
    "# `relay.qnn.quantize`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f65b6250",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import tvm\n",
    "from tvm import te\n",
    "import numpy as np\n",
    "from tvm import relay\n",
    "from tvm.contrib import graph_executor\n",
    "from tvm.relay.testing import run_infer_type\n",
    "\n",
    "\n",
    "def quantize_test_driver(in_dtype, quant_args, axis, out_dtype, in_data, verify_output_data):\n",
    "    \"\"\"\n",
    "    QNN量化测试驱动函数\n",
    "\n",
    "    该函数用于执行量化操作并验证结果正确性。它创建一个包含量化操作的\n",
    "    Relay计算图，编译并执行，然后将结果与预期输出进行比较。\n",
    "\n",
    "    参数:\n",
    "    in_dtype (str): 输入数据类型\n",
    "    quant_args (dict): 量化参数，包含'out_zero_point'和'out_scale'\n",
    "    axis (int): 量化轴\n",
    "    out_dtype (str): 输出数据类型\n",
    "    in_data (numpy.ndarray): 输入数据\n",
    "    verify_output_data (numpy.ndarray): 预期输出数据\n",
    "\n",
    "    返回:\n",
    "    None\n",
    "    \"\"\"\n",
    "    # 获取输入数据形状\n",
    "    shape = in_data.shape\n",
    "    # 创建Relay输入变量\n",
    "    input_data = relay.var(\"input_data\", shape=shape, dtype=in_dtype)\n",
    "    # 创建量化参数常量\n",
    "    output_zero_point = relay.const(quant_args[\"out_zero_point\"])\n",
    "    output_scale = relay.const(quant_args[\"out_scale\"])\n",
    "    # 创建量化操作\n",
    "    quantized_output = relay.qnn.quantize(\n",
    "        input_data,\n",
    "        output_scale=output_scale,\n",
    "        output_zero_point=output_zero_point,\n",
    "        axis=axis,\n",
    "        out_dtype=out_dtype,\n",
    "    )\n",
    "    # 构建Relay函数和IR模块\n",
    "    mod = relay.Function(relay.analysis.free_vars(quantized_output), quantized_output)\n",
    "    mod = tvm.IRModule.from_expr(mod)\n",
    "    # 编译模块\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        graph, lib, params = relay.build(mod, \"llvm\", params=None)\n",
    "        # 创建运行时模块\n",
    "        rt_mod = graph_executor.create(graph, lib, device=tvm.cpu(0))\n",
    "        # 设置输入数据\n",
    "        rt_mod.set_input(input_data=in_data)\n",
    "        rt_mod.set_input(**params)\n",
    "        # 执行计算\n",
    "        rt_mod.run()\n",
    "        # 获取输出结果\n",
    "        res = rt_mod.get_output(0).numpy()\n",
    "        # 验证结果\n",
    "        np.testing.assert_equal(res, verify_output_data)\n",
    "        assert res.dtype == out_dtype\n",
    "\n",
    "\n",
    "def test_float32_to_uint8():\n",
    "    \"\"\"\n",
    "    测试float32到uint8的量化\n",
    "\n",
    "    该函数测试将float32类型数据量化为uint8类型数据的功能，\n",
    "    使用特定的零点和缩放因子参数。\n",
    "    \"\"\"\n",
    "    # 创建测试输入数据\n",
    "    data = (\n",
    "        np.array([-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64])\n",
    "        .astype(\"float32\")\n",
    "        .reshape((2, 5))\n",
    "    )\n",
    "    # 预期输出数据\n",
    "    output = np.array([0, 1, 2, 3, 4, 251, 252, 253, 254, 255]).astype(\"uint8\").reshape((2, 5))\n",
    "    # 量化参数：零点=127，缩放因子=0.5\n",
    "    quant_args = {\"out_zero_point\": np.int32(127), \"out_scale\": np.float32(0.5)}\n",
    "    # 调用测试驱动函数\n",
    "    quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=-1,  # 最后一个维度作为量化轴\n",
    "        out_dtype=\"uint8\",\n",
    "        in_data=data,\n",
    "        verify_output_data=output,\n",
    "    )\n",
    "\n",
    "\n",
    "def test_float32_to_int8():\n",
    "    \"\"\"\n",
    "    测试float32到int8的量化\n",
    "\n",
    "    该函数测试将float32类型数据量化为int8类型数据的功能，\n",
    "    使用特定的零点和缩放因子参数。\n",
    "    \"\"\"\n",
    "    # 创建测试输入数据\n",
    "    data = (\n",
    "        np.array([-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64])\n",
    "        .astype(\"float32\")\n",
    "        .reshape((2, 5))\n",
    "    )\n",
    "    # 预期输出数据\n",
    "    output = (\n",
    "        np.array([-128, -127, -126, -125, -124, 123, 124, 125, 126, 127])\n",
    "        .astype(\"int8\")\n",
    "        .reshape((2, 5))\n",
    "    )\n",
    "    # 量化参数：零点=-1，缩放因子=0.5\n",
    "    quant_args = {\"out_zero_point\": np.int32(-1), \"out_scale\": np.float32(0.5)}\n",
    "    # 调用测试驱动函数\n",
    "    quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=-1,  # 最后一个维度作为量化轴\n",
    "        out_dtype=\"int8\",\n",
    "        in_data=data,\n",
    "        verify_output_data=output,\n",
    "    )\n",
    "\n",
    "\n",
    "def test_float32_to_uint16():\n",
    "    \"\"\"\n",
    "    测试float32到uint16的量化\n",
    "\n",
    "    该函数测试将float32类型数据量化为uint16类型数据的功能，\n",
    "    使用特定的零点和缩放因子参数。\n",
    "    \"\"\"\n",
    "    # 创建测试输入数据\n",
    "    data = (\n",
    "        np.array([-6553, -6552.8, -6552.6, -6552.4, -6552.2, 6553.2, 6553.4, 6553.6, 6553.8, 6554])\n",
    "        .astype(\"float32\")\n",
    "        .reshape((2, 5))\n",
    "    )\n",
    "    # 预期输出数据\n",
    "    output = (\n",
    "        np.array([0, 1, 2, 3, 4, 65531, 65532, 65533, 65534, 65535])\n",
    "        .astype(\"uint16\")\n",
    "        .reshape((2, 5))\n",
    "    )\n",
    "    # 量化参数：零点=32765，缩放因子=0.2\n",
    "    quant_args = {\"out_zero_point\": np.int32(32765), \"out_scale\": np.float32(0.2)}\n",
    "    # 调用测试驱动函数\n",
    "    quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=-1,  # 最后一个维度作为量化轴\n",
    "        out_dtype=\"uint16\",\n",
    "        in_data=data,\n",
    "        verify_output_data=output,\n",
    "    )\n",
    "\n",
    "\n",
    "def test_scalar_float32_to_int8():\n",
    "    \"\"\"\n",
    "    测试标量float32到int8的量化\n",
    "\n",
    "    该函数测试将标量float32类型数据量化为int8类型数据的功能，\n",
    "    验证对标量输入的处理正确性。\n",
    "    \"\"\"\n",
    "    # 创建标量测试输入数据\n",
    "    data = np.array(-63.5).astype(\"float32\")\n",
    "    # 预期输出数据\n",
    "    output = np.array(-128).astype(\"int8\")\n",
    "    # 量化参数：零点=-1，缩放因子=0.5\n",
    "    quant_args = {\"out_zero_point\": np.int32(-1), \"out_scale\": np.float32(0.5)}\n",
    "    # 调用测试驱动函数\n",
    "    quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=-1,  # 标量的唯一维度作为量化轴\n",
    "        out_dtype=\"int8\",\n",
    "        in_data=data,\n",
    "        verify_output_data=output,\n",
    "    )\n",
    "\n",
    "\n",
    "def test_channelwise_axis_0():\n",
    "    \"\"\"\n",
    "    测试通道维度为0的通道量化\n",
    "\n",
    "    该函数测试针对多维数据，以第0维作为通道维度进行通道量化的功能，\n",
    "    每个通道使用不同的量化参数。\n",
    "    \"\"\"\n",
    "    # 创建测试输入数据\n",
    "    data = (\n",
    "        np.array([-63.5, -63, -62.5, -62, -61.5, 30, 31, 31.5, 31.75, 32])\n",
    "        .astype(\"float32\")\n",
    "        .reshape((2, 5))\n",
    "    )\n",
    "    # 预期输出数据\n",
    "    output = np.array([0, 1, 2, 3, 4, 243, 247, 249, 250, 251]).astype(\"uint8\").reshape((2, 5))\n",
    "    # 通道量化参数：每个通道有自己的零点和缩放因子\n",
    "    quant_args = {\n",
    "        \"out_zero_point\": np.array([127, 123]).astype(\"int32\"),  # 两个通道的零点\n",
    "        \"out_scale\": np.array([0.5, 0.25]).astype(\"float32\"),   # 两个通道的缩放因子\n",
    "    }\n",
    "\n",
    "    # 调用测试驱动函数，axis=0表示按第0维（通道维）进行量化\n",
    "    quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=0,\n",
    "        out_dtype=\"uint8\",\n",
    "        in_data=data,\n",
    "        verify_output_data=output,\n",
    "    )\n",
    "\n",
    "\n",
    "def test_channelwise_axis_1():\n",
    "    \"\"\"\n",
    "    测试通道维度为1的通道量化\n",
    "\n",
    "    该函数测试针对多维数据，以第1维作为通道维度进行通道量化的功能，\n",
    "    每个通道使用不同的量化参数。\n",
    "    \"\"\"\n",
    "    # 创建测试输入数据并转置，使通道维变为第1维\n",
    "    data = np.transpose(\n",
    "        np.array([-63.5, -63, -62.5, -62, -61.5, 30, 31, 31.5, 31.75, 32])\n",
    "        .astype(\"float32\")\n",
    "        .reshape((2, 5))\n",
    "    )\n",
    "    # 预期输出数据并转置\n",
    "    output = np.transpose(\n",
    "        np.array([0, 1, 2, 3, 4, 243, 247, 249, 250, 251]).astype(\"uint8\").reshape((2, 5))\n",
    "    )\n",
    "    # 通道量化参数：每个通道有自己的零点和缩放因子\n",
    "    quant_args = {\n",
    "        \"out_zero_point\": np.array([127, 123]).astype(\"int32\"),  # 两个通道的零点\n",
    "        \"out_scale\": np.array([0.5, 0.25]).astype(\"float32\"),   # 两个通道的缩放因子\n",
    "    }\n",
    "\n",
    "    # 调用测试驱动函数，axis=-1表示按最后一维（通道维）进行量化\n",
    "    quantize_test_driver(\n",
    "        in_dtype=\"float32\",\n",
    "        quant_args=quant_args,\n",
    "        axis=-1,\n",
    "        out_dtype=\"uint8\",\n",
    "        in_data=data,\n",
    "        verify_output_data=output,\n",
    "    )\n",
    "\n",
    "\n",
    "def test_dynamic_quantize():\n",
    "    \"\"\"\n",
    "    测试动态量化功能\n",
    "\n",
    "    该函数测试动态量化功能，其中量化参数（缩放因子和零点）是动态变量，\n",
    "    而不是固定常量。这验证了量化操作在运行时接受动态输入的能力。\n",
    "    \"\"\"\n",
    "    # 创建动态变量\n",
    "    x = relay.var(\"x\", shape=(1, 2, 3, 4), dtype=\"float32\")  # 输入数据变量\n",
    "    scale_var = relay.var(\"scale\", shape=(), dtype=\"float32\")  # 缩放因子变量\n",
    "    zp_var = relay.var(\"zp\", shape=(), dtype=\"int32\")  # 零点变量\n",
    "\n",
    "    # 创建量化操作，使用动态计算的缩放因子和零点\n",
    "    q_x = relay.qnn.quantize(x, scale_var * scale_var, zp_var + zp_var)\n",
    "    # 推断类型\n",
    "    tt = run_infer_type(q_x)\n",
    "\n",
    "    # 验证输出类型\n",
    "    assert tt.checked_type == relay.TensorType((1, 2, 3, 4), \"int8\")\n",
    "    # 创建Relay函数\n",
    "    func = relay.Function([x, scale_var, zp_var], q_x)\n",
    "    # 生成测试数据\n",
    "    data = np.random.uniform(size=(1, 2, 3, 4)).astype(\"float32\")\n",
    "    scale = np.array(1).astype(\"float32\")\n",
    "    zp = np.array(0).astype(\"int32\")\n",
    "\n",
    "    # 构建IR模块\n",
    "    mod = tvm.ir.IRModule.from_expr(func)\n",
    "\n",
    "    # 在可用目标上编译\n",
    "    for target, dev in tvm.testing.enabled_targets():\n",
    "        # 禁用AlterOpLayout优化（待修复）\n",
    "        with relay.build_config(opt_level=3, disabled_pass=[\"AlterOpLayout\"]):\n",
    "            lib = relay.build(mod, target=target)\n",
    "\n",
    "    # 创建运行时模块并执行\n",
    "    module = graph_executor.GraphModule(lib[\"default\"](dev))\n",
    "    module.set_input(**{\"x\": data, \"scale\": scale, \"zp\": zp})\n",
    "    module.run()\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # 执行所有测试函数\n",
    "    test_float32_to_uint8()  # 测试float32到uint8的量化\n",
    "    test_float32_to_int8()   # 测试float32到int8的量化\n",
    "    test_float32_to_uint16() # 测试float32到uint16的量化\n",
    "    test_scalar_float32_to_int8()  # 测试标量float32到int8的量化\n",
    "    test_channelwise_axis_0()  # 测试通道维度为0的通道量化\n",
    "    test_channelwise_axis_1()  # 测试通道维度为1的通道量化\n",
    "    test_dynamic_quantize()  # 测试动态量化功能\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
